import h5py
import numpy as np
import cv2
import os
from tqdm import tqdm
import argparse
import math
import matplotlib.pyplot as plt
import matplotlib

matplotlib.use("Agg")


"""
Load /home/kisara/data/Fnii_data/picking_banana_hdf5/20250515/episode_0.hdf5
Dataset: action, shape: (212, 14), dtype: float32
Group: observations
Group: observations/images
Dataset: observations/images/cam1, shape: (212, 480, 640, 3), dtype: uint8
Dataset: observations/images/cam2, shape: (212, 480, 640, 3), dtype: uint8
Dataset: observations/images/cam3, shape: (212, 480, 640, 3), dtype: uint8
Dataset: observations/qpos, shape: (212, 14), dtype: float32
"""


def list_hdf5_contents(file_path):
    print(f"\nList the contents of the HDF5 file: {file_path}")

    def print_dataset(name, obj):
        if isinstance(obj, h5py.Dataset):
            print(f"- Dataset: {name}, Shape: {obj.shape}, Dtype: {obj.dtype}")
        elif isinstance(obj, h5py.Group):
            print(f"Group: {name}")

    with h5py.File(file_path, "r") as f:
        print(f"File attributes: {dict(f.attrs)}")
        print("\nAll datasets/groups:")
        print(f["observations/qpos"][:1])
        f.visititems(print_dataset)


def hdf5_to_all_videos(
    file_path,
    output_dir,
    base_filename="video",
    target_height=720,
    max_cols=4,
    fps=30,
):
    """
    将HDF5文件中的所有图像数据集转换为单独的视频和网格拼接视频

    参数:
        file_path (str): 输入HDF5文件路径
        output_dir (str): 输出视频目录
        base_filename (str): 输出文件基础名称
        target_height (int): 目标高度用于调整图像大小 (默认720)
        max_cols (int): 网格最大列数 (默认4)
        fps (int): 输出视频帧率 (默认30)
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    print(f"Processing HDF5 file: {file_path}")

    with h5py.File(file_path, "r") as f:
        # 列出文件中所有数据集
        print("Datasets in file:")
        for name, obj in f.items():
            if isinstance(obj, h5py.Dataset):
                print(f"  - {name}: shape={obj.shape}, dtype={obj.dtype}")
            elif isinstance(obj, h5py.Group):
                print(f"  - Group: {name} (contains {len(obj.keys())} items)")

        # 识别图像数据集
        image_datasets = find_image_datasets_in_object(f)

        if not image_datasets:
            print("No valid image datasets found in the file")
            return

        print(f"Found {len(image_datasets)} valid image datasets:")

        # 确定所有数据集的共同帧数
        min_frames = min(f[ds].shape[0] for ds in image_datasets)
        print(f"Common frame count across all datasets: {min_frames}")

        # 创建网格视频
        grid_output = os.path.join(output_dir, f"{base_filename}_grid.mp4")
        print(f"\nCreating grid video: {grid_output}")
        create_grid_video(
            f, image_datasets, min_frames, grid_output, target_height, max_cols, fps
        )

        print("\nAll processing completed!")


def find_image_datasets_in_object(f):
    """在打开的HDF5文件对象中查找图像数据集"""
    image_datasets = []

    def find_datasets(name, obj):
        # 如果是数据集且维度≥3
        if isinstance(obj, h5py.Dataset) and len(obj.shape) >= 3:
            # 检查是否为图像数据集（通道数为1, 3, 4或只有3维）
            if (obj.shape[-1] in [1, 3, 4] or len(obj.shape) == 3) and obj.dtype in (
                np.uint8,
                np.float32,
            ):
                # 排除小数据集（可能不是图像）
                if obj.size > 10000:
                    image_datasets.append(name)

    # 遍历所有项目
    f.visititems(find_datasets)

    return image_datasets


def convert_single_dataset(dataset, output_video, fps=30):
    """将单个图像数据集转换为视频文件"""
    # 检查数据集维度
    if len(dataset.shape) < 3:
        print(
            f"Invalid dataset dimensions: {dataset.shape} (expected at least 3 dimensions)"
        )
        return

    # 获取维度信息
    height, width = dataset.shape[1], dataset.shape[2]

    # 创建视频写入器
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter(output_video, fourcc, fps, (width, height))

    if not out.isOpened():
        print(f"Failed to create video writer for {output_video}")
        return

    print(f"  Creating video: {width}x{height}, {fps} FPS")

    # 处理每一帧
    for i in tqdm(range(dataset.shape[0]), desc="Processing frames"):
        frame = dataset[i]

        # 转换颜色格式
        if frame.shape[-1] == 3:  # RGB
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        elif frame.shape[-1] == 4:  # RGBA
            frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR)
        elif len(frame.shape) == 2:  # 灰度
            frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)

        out.write(frame)

    out.release()
    print(f"  Video saved: {output_video}")


def create_grid_video(
    f, dataset_names, min_frames, output_video, target_height=720, max_cols=4, fps=30
):
    """创建网格布局视频，将所有图像数据集拼接在一起"""
    # 确定网格布局
    n_cols = min(len(dataset_names), max_cols)
    n_rows = math.ceil(len(dataset_names) / n_cols)

    print(f"Grid layout: {n_rows} rows x {n_cols} columns")

    # 计算网格单元大小（所有图像调整到目标高度）
    cell_widths = []
    for ds_name in dataset_names:
        ds = f[ds_name]
        if len(ds.shape) >= 3:
            aspect_ratio = ds.shape[2] / ds.shape[1]
            width = int(target_height * aspect_ratio)
            cell_widths.append(width)
        else:
            cell_widths.append(target_height)  # 备选

    max_width = max(cell_widths) if cell_widths else target_height

    grid_width = n_cols * max_width
    grid_height = n_rows * target_height

    print(f"Grid dimensions: {grid_width} x {grid_height}")

    # 创建视频写入器
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter(output_video, fourcc, fps, (grid_width, grid_height))

    if not out.isOpened():
        print(f"Failed to create grid video writer for {output_video}")
        return

    # 处理每一帧
    for i in tqdm(range(min_frames), desc="Creating grid video"):
        # 创建新的网格帧
        grid_frame = np.zeros((grid_height, grid_width, 3), dtype=np.uint8)

        # 用所有数据集填充网格
        for idx, ds_name in enumerate(dataset_names):
            ds = f[ds_name]

            # 如果帧索引超出范围则跳过
            if i >= ds.shape[0]:
                continue

            # 获取当前帧
            frame = ds[i]

            # 处理颜色格式
            if len(frame.shape) == 2:  # 灰度
                frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
            elif frame.shape[-1] == 3:  # RGB
                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            elif frame.shape[-1] == 4:  # RGBA
                frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR)

            # 调整大小保持宽高比
            if frame.shape[0] > 0:  # 避免除以0
                aspect_ratio = frame.shape[1] / frame.shape[0]
                new_width = int(target_height * aspect_ratio)
                resized_frame = cv2.resize(frame, (new_width, target_height))
            else:
                new_width = max_width
                resized_frame = np.zeros((target_height, max_width, 3), dtype=np.uint8)

            # 计算网格位置
            row = idx // n_cols
            col = idx % n_cols

            # 计算起始位置
            x_start = col * max_width
            y_start = row * target_height

            # 放置图像到网格
            if new_width <= max_width:
                # 图像宽度小于等于网格单元宽度
                grid_frame[
                    y_start : y_start + target_height, x_start : x_start + new_width
                ] = resized_frame
            else:
                # 图像宽度大于网格单元宽度 - 居中对齐
                x_offset = (max_width - new_width) // 2
                grid_frame[
                    y_start : y_start + target_height,
                    x_start + x_offset : x_start + x_offset + new_width,
                ] = resized_frame

        # 添加帧号信息
        cv2.putText(
            grid_frame,
            f"Frame: {i+1}/{min_frames}",
            (20, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            1,
            (0, 255, 255),
            2,
        )

        # 添加数据集名称
        for idx, ds_name in enumerate(dataset_names):
            row = idx // n_cols
            col = idx % n_cols

            # 简化名称用于显示
            display_name = ds_name.split("/")[-1]  # 使用最后一部分
            if len(display_name) > 20:
                display_name = display_name[:17] + "..."

            # 计算位置
            x_pos = col * max_width + 10
            y_pos = (row + 1) * target_height - 10

            # 添加文本
            cv2.putText(
                grid_frame,
                display_name,
                (x_pos, y_pos),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.6,
                (255, 255, 255),
                1,
            )

        # 将网格帧写入视频
        out.write(grid_frame)

    out.release()
    print(f"Grid video saved: {output_video}")


def hdf5_save_arm(file_path, output_dir):
    has_gpos = False
    with h5py.File(file_path, "r") as f:
        # 读取 datasets
        action_qpos = f["action"][:]
        obs_qpos = f["observations/qpos"][:]
        if "action_gpos" in f:
            has_gpos = True
            action_gpos = f["action_gpos"][:]
            obs_gpos = f["observations/gpos"][:]

    # 时间步
    timesteps = np.arange(action_qpos.shape[0])
    action_dim = action_qpos.shape[1]
    print(f"action shape: {action_qpos.shape}, obs shape: {obs_qpos.shape}")

    # -------- gpos 对照图 --------
    if has_gpos:
        rows = int(np.ceil(action_dim / 2))  # 每行放2个子图
        cols = 2
        fig, axes = plt.subplots(rows, cols, figsize=(12, rows * 3), squeeze=False)

        for i in range(action_dim):
            r, c = divmod(i, cols)
            ax = axes[r, c]
            ax.plot(timesteps, action_gpos[:, i], label="action_gpos")
            ax.plot(timesteps, obs_gpos[:, i], linestyle="--", label="obs_gpos")
            ax.set_title(f"dim {i}")
            ax.set_xlabel("Timestep")
            ax.set_ylabel("Value")
            ax.grid(True)
            ax.legend()

        plt.suptitle("action_gpos vs obs_gpos")
        plt.tight_layout(rect=[0, 0, 1, 0.96])
        plt.savefig(os.path.join(output_dir, "gpos_compare.png"))
        plt.close()

    # -------- qpos 对照图 --------
    rows = int(np.ceil(action_dim / 2))  # 每行放2个子图
    cols = 2
    fig, axes = plt.subplots(rows, cols, figsize=(12, rows * 3), squeeze=False)

    for i in range(action_dim):
        r, c = divmod(i, cols)
        ax = axes[r, c]
        ax.plot(timesteps, action_qpos[:, i], label="action_qpos")
        ax.plot(timesteps, obs_qpos[:, i], linestyle="--", label="obs_qpos")
        ax.set_title(f"dim {i}")
        ax.set_xlabel("Timestep")
        ax.set_ylabel("Value")
        ax.grid(True)
        ax.legend()

    plt.suptitle("action_qpos vs obs_qpos")
    plt.tight_layout(rect=[0, 0, 1, 0.96])
    plt.savefig(os.path.join(output_dir, "qpos_compare.png"))
    plt.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input",
        type=str,
        help="hdf5 file path",
        default="/home/kisara/data/Fnii_vla/data/open_roaster_0612",
    )
    parser.add_argument(
        "--list", action="store_true", help="list the contents of the file"
    )
    parser.add_argument("--episode_name", type=str, default=None, help="episode name")
    parser.add_argument("--video", type=str, help="mp4 save path")
    parser.add_argument("--check_all", type=str, help="save video and joint")
    parser.add_argument("--fps", type=int, default=20)

    args = parser.parse_args()

    # 执行相应操作
    if args.list:
        file_path = args.input
        if file_path[-1] != "/":
            file_path = file_path + "/"
        file_path = file_path + args.episode_name
        list_hdf5_contents(file_path)
    elif args.video:
        file_path = args.input
        if file_path[-1] != "/":
            file_path = file_path + "/"
        file_path = file_path + args.episode_name
        hdf5_to_all_videos(
            file_path=file_path,  # 替换为您的HDF5文件路径
            output_dir=args.video,  # 输出目录
            # target_height=360,  # 目标高度
            max_cols=3,  # 最大列数
            fps=args.fps,  # 帧率
        )
    elif args.check_all:
        file_path = args.input
        if file_path[-1] != "/":
            file_path = file_path + "/"
        file_path = file_path + args.episode_name
        hdf5_to_all_videos(
            file_path=file_path,  # 替换为您的HDF5文件路径
            output_dir=args.check_all,  # 输出目录
            # target_height=360,  # 目标高度
            max_cols=3,  # 最大列数
            fps=args.fps,  # 帧率
        )
        hdf5_save_arm(file_path=file_path, output_dir=args.check_all)
    else:
        print("Please use --list, --find-images, --extract-images or --video")
